package com.pt25.base.util;

import jakarta.websocket.*;
import lombok.*;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.pool2.PooledObject;
import org.apache.commons.pool2.PooledObjectFactory;
import org.apache.commons.pool2.impl.DefaultPooledObject;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.apache.commons.pool2.impl.GenericObjectPoolConfig;

import java.io.IOException;
import java.net.URI;
import java.nio.ByteBuffer;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Slf4j
public class WsClientUtil1 {

    private static final Map<String, WebSocketClientPool> webSocketClientPoolMap = new HashMap<>();

    public static WebSocketClientPool getWebsocketClientPool(WSConfig wsConfig) {
        String mapKey = wsConfig.getWebSocketUrl().split("\\?")[0];
        if (null == webSocketClientPoolMap.get(mapKey)) {
            URI endpointURI = URI.create(wsConfig.getWebSocketUrl());
            GenericObjectPoolConfig<WebSocketClientEndpoint> config = new GenericObjectPoolConfig<>();
            config.setMaxTotal(wsConfig.getMaxTotal());
            config.setMinIdle(wsConfig.getMinIdle());
            config.setMaxIdle(wsConfig.getMaxIdle());
            config.setTestOnBorrow(wsConfig.getTestOnBorrow());
            config.setTestOnReturn(wsConfig.getTestOnReturn());
            WebSocketClientPool pool = new WebSocketClientPool(config, endpointURI, wsConfig.getHeaders());
            webSocketClientPoolMap.put(mapKey, pool);
        }
        return webSocketClientPoolMap.get(mapKey);
    }

    @NoArgsConstructor
    @AllArgsConstructor
    @Builder
    @Data
    public static class WSConfig {
        private String webSocketUrl;
        @Builder.Default
        private Map<String, String> headers = new HashMap<>();
        @Builder.Default
        private Integer maxTotal = 20;
        @Builder.Default
        private Integer minIdle = 1;
        @Builder.Default
        private Integer maxIdle = 5;
        @Builder.Default
        private Boolean testOnBorrow = true;
        @Builder.Default
        private Boolean testOnReturn = true;
        @Builder.Default
        private Long timeBetweenEvictionRunsMillis = 30000L;//30秒
        @Builder.Default
        private Long minEvictableIdleTimeMillis = 60000L;//60秒
    }

    public static class WebSocketClientEndpoint {

        private Session userSession;
        @Setter
        private MessageHandler messageHandler;

        public WebSocketClientEndpoint(URI endpointURI, Map<String, String> headers) {
            try {
                WebSocketContainer container = ContainerProvider.getWebSocketContainer();
                if (MapUtil.isNotEmpty(headers)) {
                    ClientEndpointConfig clientConfig = ClientEndpointConfig.Builder.create()
                            .configurator(new CustomConfigurator(headers))
                            .build();

                    container.connectToServer(new Endpoint() {
                        @Override
                        public void onOpen(Session session, EndpointConfig config) {
                            WebSocketClientEndpoint.this.userSession = session;
                            System.out.println("WebSocket connection opened");

                            session.addMessageHandler(new jakarta.websocket.MessageHandler.Whole<String>() {
                                @Override
                                public void onMessage(String message) {
                                    if (messageHandler != null) {
                                        messageHandler.handleMessage(message);
                                    }
                                }
                            });

                            session.addMessageHandler(new jakarta.websocket.MessageHandler.Whole<ByteBuffer>() {
                                @Override
                                public void onMessage(ByteBuffer bytes) {
                                    if (messageHandler != null) {
                                        messageHandler.handleMessage(bytes);
                                    }
                                }
                            });
                        }

                        @Override
                        public void onClose(Session session, CloseReason closeReason) {
                            System.out.println("WebSocket connection closed: " + closeReason);
                            WebSocketClientEndpoint.this.userSession = null;
                            if (messageHandler != null) {
                                messageHandler.handleClose(session, closeReason);
                            }
                        }

                        @Override
                        public void onError(Session session, Throwable thr) {
                            thr.printStackTrace();
                            if (messageHandler != null) {
                                messageHandler.handleError(session, thr);
                            }
                        }
                    }, clientConfig, endpointURI);
                } else {
                    container.connectToServer(this, endpointURI);
                }
            } catch (DeploymentException | IOException e) {
                throw new RuntimeException(e);
            }
        }

        public void sendMessage(String message) {
            if (this.userSession != null) {
                this.userSession.getAsyncRemote().sendText(message);
            }
        }

        public void sendMessage(ByteBuffer message) {
            if (this.userSession != null) {
                this.userSession.getAsyncRemote().sendBinary(message);
            }
        }

        public void close() {
            try {
                if (this.userSession != null) {
                    this.userSession.close();
                }
            } catch (IOException e) {
                log.error(e.getMessage(), e);
            }
        }

        public boolean isOpen() {
            return this.userSession != null && this.userSession.isOpen();
        }

        public interface MessageHandler {
            void handleMessage(String message);

            void handleMessage(ByteBuffer bytes);

            void handleError(Session session, Throwable thr);

            void handleClose(Session session, CloseReason closeReason);
        }

        private static class CustomConfigurator extends ClientEndpointConfig.Configurator {
            private final Map<String, String> headers;

            public CustomConfigurator(Map<String, String> headers) {
                this.headers = headers;
            }

            @Override
            public void beforeRequest(Map<String, List<String>> headersMap) {
                headers.forEach((key, value) -> headersMap.put(key, Collections.singletonList(value)));
            }
        }
    }

    public static class WebSocketClientPool {
        private final GenericObjectPool<WebSocketClientEndpoint> pool;

        public WebSocketClientPool(GenericObjectPoolConfig<WebSocketClientEndpoint> config, URI endpointURI, Map<String, String> headers) {
            this.pool = new GenericObjectPool<>(new WebSocketClientFactory(endpointURI, headers), config);
        }

        public WebSocketClientEndpoint borrowClient() throws Exception {
            return pool.borrowObject();
        }

        public void returnClient(WebSocketClientEndpoint client) {
            pool.returnObject(client);
        }

        public void close() {
            pool.close();
        }

        private record WebSocketClientFactory(URI endpointURI,
                                              Map<String, String> headers) implements PooledObjectFactory<WebSocketClientEndpoint> {

            @Override
                    public PooledObject<WebSocketClientEndpoint> makeObject() throws Exception {
                        WebSocketClientEndpoint client = new WebSocketClientEndpoint(endpointURI, headers);
                        return new DefaultPooledObject<>(client);
                    }

                    @Override
                    public void destroyObject(PooledObject<WebSocketClientEndpoint> pooledObject) throws Exception {
                        WebSocketClientEndpoint client = pooledObject.getObject();
                        client.close();
                    }

                    @Override
                    public boolean validateObject(PooledObject<WebSocketClientEndpoint> pooledObject) {
                        WebSocketClientEndpoint client = pooledObject.getObject();
                        return client.isOpen();
                    }

                    @Override
                    public void activateObject(PooledObject<WebSocketClientEndpoint> pooledObject) throws Exception {
                        // No activation needed
                    }

                    @Override
                    public void passivateObject(PooledObject<WebSocketClientEndpoint> pooledObject) throws Exception {
                        // No passivation needed
                    }
                }
    }
}
